package com.yjj.demo.common.interceptor;

import com.yjj.demo.common.annotation.QPSLimiter;
import lombok.extern.log4j.Log4j2;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 *
 *
 * @author yangjiajia
 * @createdAt 2023/1/19 14:27
 */
@Log4j2
public class QPSLimiterInterceptor implements HandlerInterceptor {

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    private static final String CONTENT_APPLICATION_JSON_UTF8 = "application/json;charset=UTF-8";

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            QPSLimiter qpsLimiter = handlerMethod.getMethodAnnotation(QPSLimiter.class);
            if (qpsLimiter == null) {
                return true;
            }

            String userId = request.getParameter("userId");
            String key = String.format("api_qps_%s", userId);
            int currentCount = 1;
            if (Boolean.FALSE.equals(stringRedisTemplate.hasKey(key))) {
                stringRedisTemplate.opsForValue().set(key, "1", 1, TimeUnit.SECONDS);
            } else {
                currentCount = Objects.requireNonNull(stringRedisTemplate.opsForValue().increment(key)).intValue();
            }
            if (currentCount >= qpsLimiter.count()) {
                log.error("url:{},超出qps上限,userId:{}", request.getRequestURI(), userId);
                response.setContentType(CONTENT_APPLICATION_JSON_UTF8);
                response.getWriter().print(qpsLimiter.msg());
                return false;
            }

        }
        return true;
    }
}
